import tracemalloc

from elasticsearch import Elasticsearch
import db
from pcap import Pcap


def process_file(pcap, pos, item_pos, length):
    pcap.ref()
    try:
        packet = pcap.read_packet(pos)
        if packet is None:
            print(f"couldn't read packet at {pos} packet # {item_pos} of {length}")
            pcap.unref()
            return None, None, msg1
        else:
            pcap.unref()
            return pcap, packet, None
    except Exception as e:
        pcap.unref()
        return None, None, e


def process_session_id_disk(node,packetPos, es):
    has_handle_header = False
    fileNum = 0
    item_pos = 0
    header=None

    b = bytearray(0xfffe)
    nextPacket = 0
    b_offset = 0
    packets = {}
    def pCb(pca, buffer, i):
        nonlocal b_offset, b, packets, header, nextPacket
        print(f"packet:{i}  {type(buffer)}  {len(buffer)}")
        packets[i] = buffer
        while nextPacket in packets:
            buffer = packets[nextPacket]
            del packets[nextPacket]
            nextPacket = nextPacket + 1
            if b_offset + len(buffer) > len(b):
                res = header + b[:b_offset]
                b_offset = 0
                b = bytearray(0xfffe)
            b[b_offset:b_offset + len(buffer)] = buffer
            b_offset += len(buffer)


    try:
        for pos in session['packetPos']:
            if pos < 0:
                fileNum = pos * -1
                continue
            opcap = Pcap.get(f"{node}:{fileNum}")
            if opcap.is_corrupt():
                print("文件被占用")
                continue
            elif not opcap.is_open():
                file = db.Db.get_file_by_file_id(node=node, num=fileNum, es=es)
                if not file:
                    msg = f"Only have SPI data, PCAP file no longer available. Couldn't look up {node}-{fileNum} in files index"
                    print(msg)
                    continue
                try:
                    ipcap = Pcap.get(f"{node}:{fileNum}")
                    ipcap.open(file)
                except Exception as err:
                    print(f"打开文件失败{err}")
                    continue
                if not has_handle_header:
                    header = ipcap.read_header()
                    has_handle_header = True
                pcap, packet, err_msg = process_file(ipcap, pos, item_pos, packetPos)
                if err_msg:
                    print(err_msg)
                    continue
                if pcap and packet:
                    pCb(pcap, packet, item_pos)
                    item_pos = item_pos + 1
            else:
                if not has_handle_header:
                    header = opcap.read_header()
                    has_handle_header = True
                pcap, packet, err_msg = process_file(opcap, pos, item_pos, packetPos)
                if err_msg:
                    print(err_msg)
                    continue
                if pcap and packet:
                    pCb(pcap, packet, item_pos)
                    item_pos = item_pos + 1
        return header + b[:b_offset]
    except Exception as e:
        print(e)
        return None


def fix_pos(node, pos, es):
    if pos is None or len(pos) == 0:
        return
    file = db.Db.get_file_by_file_id(node=node, num=abs(pos[0]), es=es)
    if file and "packetPosEncoding" in file:
        if file['packetPosEncoding'] == "gap0":
            last = 0
            lastgap = 0
            for i, pos_item in enumerate(pos):
                if pos[i] < 0:
                    last = 0
                else:
                    if pos[i] == 0:
                        pos[i] = last + lastgap
                    else:
                        lastgap = pos[i]
                        pos[i] += last
                    last = pos[i]
        elif file['packetPosEncoding'] == "localIndex":
            # this localPcapIndex=false in config.ini
            # "localPcapIndex = arkime_config_boolean(NULL, "localPcapIndex", FALSE);"
            return
    else:
        return


def main():
    es = Elasticsearch('http://127.0.0.1:9200')
    res = es.search(body={
        "query": {
            "match": {
                "_id": "240806-2xgUoBgyv71DdrSKg627aeXP"
            }
        },
        "size": 1
    })
    session = res['hits']['hits'][0]
    if len(session['_source']['packetPos']) == 0:
        print("no pos")
        exit(0)
    pos = session['_source']['packetPos']
    fix_pos(session['_source']['node'], pos, es)
    result=process_session_id_disk(node=session['_source']['node'],
                            packetPos=session['_source']['packetPos'],
                            es=es)
if __name__ == '__main__':
    main()
